"""
提供用于图编译的工具
"""
import json
from langchain_core.messages import ToolMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langchain_openai import ChatOpenAI
from typing import Callable

from owl_ai.domain.agent_config_entity import AgentConfigEntity
from owl_ai.service.agent_flow_service import State


class BasicToolNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self, tools: list) -> None:
        self.tools_by_name = {tool.name: tool for tool in tools}

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        outputs = []
        for tool_call in message.tool_calls:
            tool_result = self.tools_by_name[tool_call["name"]].invoke(
                tool_call["args"]
            )
            outputs.append(
                ToolMessage(
                    content=json.dumps(tool_result),
                    name=tool_call["name"],
                    tool_call_id=tool_call["id"],
                )
            )
        return {"messages": outputs}


class ChatNode(Callable):
    def __init__(self, chat_llm: ChatOpenAI = None, system_prompt: str = None):
        self.chat_llm = chat_llm
        self.system_prompt = system_prompt

    def __call__(self, state: State, config: RunnableConfig):
        messages = state.get("messages")
        llm_messages = [
            SystemMessage(content=self.system_prompt)
        ]
        llm_messages.extend(messages)
        ai_message = self.chat_llm.invoke(llm_messages)
        return {
            "messages": [
                ai_message
            ]
        }


class GraphCompile:
    @classmethod
    def compile(cls, config: AgentConfigEntity):
        """
        编译图
        """
        pass

    @staticmethod
    def chat_llm_generate(llm_config: dict, tools: list = None):
        base_url = llm_config.get("url")
        model_name = llm_config.get("modelName")
        chat_openai = ChatOpenAI(base_url=base_url, model_name=model_name, api_key="ollama")
        if tools:
            chat_openai.bind_tools(tools)
        return chat_openai


class MultiAgentGraphCompile(GraphCompile):
    """
    多智能体图编译
    """

    @classmethod
    def compile(cls, config: AgentConfigEntity):
        pass


class AppGraphCompile(GraphCompile):
    """
    应用图编译
    """

    @classmethod
    def compile(cls, config: AgentConfigEntity):
        pass
